{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### This notebook shows how to read the fastMRI dataset and apply some simple transformations to the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Testing if integration works"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import h5py\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The fastMRI dataset is distributed as a set of HDF5 files and can be read with the h5py package. Here, we show how to open a file from the multi-coil dataset. Each file corresponds to one MRI scan and contains the k-space data, ground truth and some meta data related to the scan."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_name = 'multicoil_train/file1000167.h5'\n",
    "hf = h5py.File(file_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print('Keys:', list(hf.keys()))\n",
    "print('Attrs:', dict(hf.attrs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In multi-coil MRIs, k-space has the following shape:\n",
    "(number of slices, number of coils, height, width)\n",
    "\n",
    "For single-coil MRIs, k-space has the following shape:\n",
    "(number of slices, height, width)\n",
    "\n",
    "MRIs are acquired as 3D volumes, the first dimension is the number of 2D slices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "volume_kspace = hf['kspace'][()]\n",
    "print(volume_kspace.dtype)\n",
    "print(volume_kspace.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slice_kspace = volume_kspace[20] # Choosing the 20-th slice of this volume"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see what the absolute value of k-space looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_coils(data, slice_nums, cmap=None):\n",
    "    fig = plt.figure()\n",
    "    for i, num in enumerate(slice_nums):\n",
    "        plt.subplot(1, len(slice_nums), i + 1)\n",
    "        plt.imshow(data[num], cmap=cmap)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_coils(np.log(np.abs(slice_kspace) + 1e-9), [0, 5, 10])  # This shows coils 0, 5 and 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The fastMRI repo contains some utlity functions to convert k-space into image space. These functions work on PyTorch Tensors. The to_tensor function can convert Numpy arrays to PyTorch Tensors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import fastmri\n",
    "from fastmri.data import transforms as T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slice_kspace2 = T.to_tensor(slice_kspace)      # Convert from numpy array to pytorch tensor\n",
    "slice_image = fastmri.ifft2c(slice_kspace2)           # Apply Inverse Fourier Transform to get the complex image\n",
    "slice_image_abs = fastmri.complex_abs(slice_image)   # Compute absolute value to get a real image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_coils(slice_image_abs, [0, 5, 10], cmap='gray')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, each coil in a multi-coil MRI scan focusses on a different region of the image. These coils can be combined into the full image using the Root-Sum-of-Squares (RSS) transform."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slice_image_rss = fastmri.rss(slice_image_abs, dim=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(np.abs(slice_image_rss.numpy()), cmap='gray')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So far, we have been looking at fully-sampled data. We can simulate under-sampled data by creating a mask and applying it to k-space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastmri.data.subsample import RandomMaskFunc\n",
    "mask_func = RandomMaskFunc(center_fractions=[0.04], accelerations=[8])  # Create the mask function object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "masked_kspace, mask = T.apply_mask(slice_kspace2, mask_func)   # Apply the mask to k-space"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see what the subsampled image looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampled_image = fastmri.ifft2c(masked_kspace)           # Apply Inverse Fourier Transform to get the complex image\n",
    "sampled_image_abs = fastmri.complex_abs(sampled_image)   # Compute absolute value to get a real image\n",
    "sampled_image_rss = fastmri.rss(sampled_image_abs, dim=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(np.abs(sampled_image_rss.numpy()), cmap='gray')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.2 64-bit ('fastmri': conda)",
   "language": "python",
   "name": "python38264bitfastmricondae6ff12e14a6d4c7fab7686eaa5d0b897"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7-final"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}